#include "util.hpp"

namespace AVLT
{



template<class K, class V>
class Node
{
public:
    Node(const std::pair<K, V>& data)
        :data_(data), left_(nullptr), right_(nullptr),
         bf_(0), parent_(nullptr)
    {}
public:
    std::pair<K, V> data_;   // 数据
    Node<K, V>* left_;       // 右子树
    Node<K, V>* right_;      // 左子树
    Node<K, V>* parent_;     // 父节点
    int bf_;                 // 平衡因子 -1 代表左边高 1 +1 代表右边高1
};


template<class K, class V>
class AVLTree
{
public:
    AVLTree()
        :root_(nullptr)
    {}
    bool Insert(const std::pair<K, V>& data)
    {
        if(root_ == nullptr)
        {
            root_ = new Node<K, V>(data);
            root_->parent_ = nullptr;
        }
        else
        {
            // 1. 先插入
            Node<K, V>* cur = root_;
            Node<K, V>* parent = root_;
            bool isLeft = true;
            while(cur)
            {
                parent = cur;
                if(data.first < cur->data_.first)
                {
                    if(cur->left_ == nullptr)
                    {
                        cur->left_ = new Node<K, V>(data);
                        cur->left_->parent_ = cur;
                        break;
                    }
                    cur = cur->left_;
                }
                else if(data.first > cur->data_.first)
                {
                    if(cur->right_ == nullptr)
                    {
                        cur->right_ = new Node<K, V>(data);
                        cur->right_->parent_ = cur;
                        break;
                    }
                    cur = cur->right_;
                }
                else
                {
                    return false;
                }
            }
            // 2. 更新平衡因子
            while(parent)
            {
                if(cur == parent->left_)
                {
                    --parent->bf_;
                }
                else
                {
                    ++parent->bf_;

                }
                // 判断是否需要继续更新
                if(parent->bf_ == 0)
                {
                    break;
                }
                else if(parent->bf_ == 1 || parent->bf_ == -1)
                {
                    // 旋转调整
                }
                else if(parent->bf_ == 2 || parent->bf_ == 2)
                {
                    // 旋转调整
                    if(parent->bf_ == -2 && cur->bf_ == -1)
                    {
                        // 右单旋转
                        RotateR(parent);
                    }
                    else if(parent->bf_ == 2 && cur->bf_ == -1)
                    {
                        // 左单旋转
                        RotateL(parent);
                    }
                    else if(parent->bf_ == -2 && cur->bf_ == 1)
                    {
                        // 左右双旋
                        RotateLR(parent);
                    }
                    else if(parent->bf_ == 2 && cur->bf_ == -1)
                    {
                        // 右左双旋
                        RotateRL(parent);
                    }
                    break;
                }
                else
                {
                    // 代表之前没有平衡
                    return false;
                }
                parent = parent->parent_;
            }

        }
        return true;
    }
    bool At(const K& key, V* const result)
    {
        Node<K, V>* tmp = root_;
        while(tmp != nullptr)
        {
            if(key < tmp->data_.first)
            {
                tmp = tmp->left_;
            }
            else if(key > tmp->data_.first)
            {
                tmp = tmp->right_;
            }
            else
            {
                *result = tmp->data_.second;
                return true;
            }
        }
        return false;
    }
    void _PrePrint(const Node<K, V>* node)
    {
        if(node == nullptr) { return; }
        _PrePrint(node->left_);
        std::cout << "[ " << node->data_.first << ", " << node->data_.second << " ] ";
        _PrePrint(node->right_);

    }
    void PrePrint()
    {
        _PrePrint(root_);
        std::cout << std::endl;
    }
private:
    // 右单旋转
    void RotateR(Node<K, V>* parent)
    {
        Node<K, V>* subL = parent->left_;
        Node<K, V>* subLR = subL->right_;
        if(subLR != nullptr)
        {
            subLR->parent_ = parent;
        }
        if(parent == root_)
        {
            root_ = parent->left_;
            root_->parent_ = nullptr;
        }
        else
        {
            if(parent->parent_->left_ == parent)
            {
                parent->parent_->left_ = subL;
            }
            else
            {
                parent->parent_->right_ = subL;
            }
            subL->parent_ = parent->parent_;
        }
        subL->right_ = parent;
        parent->parent_ = subL;
        parent->left_ = subLR;
        parent->bf_ = 0;
        subL->bf_ = 0;
    }
    // 左单旋转
    void RotateL(Node<K, V>* parent)
    {
        Node<K, V>* subR = parent->right_;
        Node<K, V>* subRL = subR->left_;
        if(subRL != nullptr)
        {
            subRL->parent_ = parent;
        }
        if(parent == root_)
        {
            root_ = subR;
            root_->parent_ = nullptr;
        }
        else
        {
            if(parent->parent_->left_ == parent)
            {
                parent->parent_->left_ = subR;

            }
            else
            {
                parent->parent_->right_ = subR;
            }
        }
        parent->parent_ = subR;
        parent->right_ = subRL;
        subR->left_ = parent;
        subR->bf_ = 0;
        parent->bf_ = 0;
    }
    // 新节点插入较高左子树的右侧---左右：先左单旋再右单旋
    // 左右双旋
    void RotateLR(Node<K, V>* parent)
    {
        Node<K, V>* subL = parent->left_;
        Node<K, V>* subLR = subL->right_;
        // 暂存bf用于更新之后的bf 这是因为插在subLR的左右是未知的
        // 通过bf来得知哪边高 从而进行更新bf
        int bf = subLR->bf_;
        // 先对左进行左单旋
        RotateL(subL);
        // 再对右进行右单旋
        RotateR(subLR);

        if(1 == bf)    subL->bf_ = -1;
        else if(-1 == bf)   parent->bf_ = 1;
    }

    // 新节点插入较高右子树的左侧---右左：先右单旋再左单旋
    // 右左双旋
    void RotateRL(Node<K, V>* parent)
    {
        Node<K, V>* subR = parent->right_;
        Node<K, V>* subRL = subR->left_;
        // 暂存bf用于更新之后的bf 这是因为插在subLR的左右是未知的
        // 通过bf来得知哪边高 从而进行更新bf
        int bf = subRL->bf_;
        // 先对右进行右单旋R
        RotateR(subR);
        // 再对左进行左单旋
        RotateL(subRL);

        if(1 == bf)    parent->bf_ = -1;
        else if(-1 == bf)   subR->bf_ = 1;
    }
    Node<K, V>* root_;
};
};